//
// Created by lijia on 2021/1/12.
//

#include "mpu.h"
#include "i2c_bus.h"
#include "system.h"
#include "util.h"
#include "display.h"

static MpuData lastMpuFrame;
static MpuCalibrationData mpuCalib = {
		.gyro = { 0, 0, 0 },
		.magneticLowerBound = { -426, 139, -747 },
		.magneticUpperBound = { 93, 647, -173 },
		.magneticInclination = 0.f,
//		.magneticLowerBound = { 9999, 9999, 9999 },
//		.magneticUpperBound = { -9999, -9999, -9999 }
};
static BarometerData lastBarometerFrame;
static BarometerCalibration calib;
static I2C_HandleTypeDef *hwnd;
static int32_t t_fine;
static uint32_t lastUpdateTime = 0, lastMagUpdateTime = 0;
static uint8_t buf[8] = {0}, magCalibrate[3] = {0};
static int16_t aXq[5] = {0}, aYq[5] = {0}, aZq[5] = {0};
static uint8_t pXq = 0, pYq = 0, pZq = 0;

static void barometer_read_coefficients() {
	calib.t1 = i2c_read_reg_16le(hwnd, BARO_I2C_ADDR, REG_BARO_T1);
	calib.t2 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_T2);
	calib.t3 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_T3);
	calib.p1 = i2c_read_reg_16le(hwnd, BARO_I2C_ADDR, REG_BARO_P1);
	calib.p2 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P2);
	calib.p3 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P3);
	calib.p4 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P4);
	calib.p5 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P5);
	calib.p6 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P6);
	calib.p7 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P7);
	calib.p8 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P8);
	calib.p9 = i2c_read_reg_s16le(hwnd, BARO_I2C_ADDR, REG_BARO_P9);
}

static void barometer_setup() {
	i2c_write_reg(hwnd, BARO_I2C_ADDR, REG_BARO_CONFIG, VALUE_BARO_CONFIG);
	i2c_write_reg(hwnd, BARO_I2C_ADDR, REG_BARO_CONTROL, VALUE_BARO_CONTROL);
}

static void barometer_init() {
	uint8_t out = i2c_read_reg(hwnd, BARO_I2C_ADDR, REG_BARO_WHO_AM_I);
	if (out != VALUE_BARO_WHO_AM_I) {
		return;
	}
	barometer_read_coefficients();
	barometer_setup();
	system_mark_module_ready(MODULE_BAROMETER);
}

static void barometer_read_temperature() {
	uint8_t buf[3] = {0, 0, 0};
	int32_t var1, var2, adc_T;

	i2c_read(hwnd, BARO_I2C_ADDR, REG_BARO_TEMPERATURE, buf, 3);

	adc_T = (buf[0] << 16) | (buf[1] << 8) | buf[2];
	adc_T >>= 4;

	var1 = ((((adc_T >> 3) - ((int32_t) calib.t1 << 1))) *
			((int32_t) calib.t2)) >> 11;

	var2 = (((((adc_T >> 4) - ((int32_t) calib.t1)) *
			  ((adc_T >> 4) - ((int32_t) calib.t1))) >> 12) *
			((int32_t) calib.t3)) >> 14;

	t_fine = var1 + var2;

	float T = (float) ((t_fine * 5 + 128) >> 8);
	lastBarometerFrame.temperature = T / 100;
}

static void barometer_read_pressure() {
	uint8_t buf[3] = {0, 0, 0};
	int64_t var1, var2, p;
	int32_t adc_P;

	// Must be done first to get the t_fine variable set up
	barometer_read_temperature();

	i2c_read(hwnd, BARO_I2C_ADDR, REG_BARO_PRESSURE, buf, 3);

	adc_P = (buf[0] << 16) | (buf[1] << 8) | buf[2];
	adc_P >>= 4;

	var1 = ((int64_t) t_fine) - 128000;
	var2 = var1 * var1 * (int64_t) calib.p6;
	var2 = var2 + ((var1 * (int64_t) calib.p5) << 17);
	var2 = var2 + (((int64_t) calib.p4) << 35);
	var1 = ((var1 * var1 * (int64_t) calib.p3) >> 8) +
		   ((var1 * (int64_t) calib.p2) << 12);
	var1 = (((((int64_t) 1) << 47) + var1)) * ((int64_t) calib.p1) >> 33;

	if (var1 == 0) {
		return; // avoid exception caused by division by zero
	}
	p = 1048576 - adc_P;
	p = (((p << 31) - var2) * 3125) / var1;
	var1 = (((int64_t) calib.p9) * (p >> 13) * (p >> 13)) >> 25;
	var2 = (((int64_t) calib.p8) * p) >> 19;

	p = ((p + var1 + var2) >> 8) + (((int64_t) calib.p7) << 4);
	lastBarometerFrame.pressure = (float) p / 256;
}

static uint8_t mag_read_reg(uint8_t reg) {
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_ADDR, 0x80 | MAG_I2C_ADDR);
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_REGISTER, reg);
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_CTRL, 0x80 | 0x01);
//	HAL_Delay(10);
//	return i2c_read_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE_DATA);
	return i2c_read_reg(hwnd, MAG_I2C_ADDR, reg);
}

static void mag_read(uint8_t *out, uint8_t reg, uint8_t len) {
//	// slow read
//	for (uint8_t i = 0; i < len; i++) {
//		out[i] = mag_read_reg(reg + i);
//	}
	i2c_read(hwnd, MAG_I2C_ADDR, reg, out, len);
}

static void mag_write_reg(uint8_t reg, uint8_t data) {
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_ADDR, MAG_I2C_ADDR);
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_REGISTER, reg);
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_OUT_DATA, data);
//	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SLAVE0_CTRL, 0x81);
//	HAL_Delay(50);
	i2c_write_reg(hwnd, MAG_I2C_ADDR, reg, data);
}

static void motion_tick();

static void motion_calibrate() {
	for (uint8_t i = 0; i < 32; i++) {
		i2c_read(hwnd, MPU_I2C_ADDR, REG_GYRO_DATA, buf, 6);
		mpuCalib.gyro[X_AXIS] = (5 * s16be(buf) + 5 * mpuCalib.gyro[X_AXIS]) / 10;
		mpuCalib.gyro[Y_AXIS] = (5 * s16be(buf + 2) + 5 * mpuCalib.gyro[Y_AXIS]) / 10;
		mpuCalib.gyro[Z_AXIS] = (5 * s16be(buf + 4) + 5 * mpuCalib.gyro[Z_AXIS]) / 10;
		HAL_Delay(50);
	}

#if HAS_MAG_CALIBRATION
	display_clear();
	display_printf("Turn the device");

	uint32_t startCalibrationTime = HAL_GetTick();
	uint32_t now = HAL_GetTick();
	while (now - startCalibrationTime < 15000) {
		motion_tick();
		for (uint8_t i = 0; i < 3; i++) {
			if (lastMpuFrame.magneticScaling[i] < mpuCalib.magneticLowerBound[i]) {
				mpuCalib.magneticLowerBound[i] = lastMpuFrame.magneticScaling[i];
			}
			if (lastMpuFrame.magneticScaling[i] > mpuCalib.magneticUpperBound[i]) {
				mpuCalib.magneticUpperBound[i] = lastMpuFrame.magneticScaling[i];
			}
		}
		now = HAL_GetTick();
	}
#endif

	/*
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS, mpuCalib.gyro[0] >> 8);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS + 1, mpuCalib.gyro[0] & 0xff);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS + 2, mpuCalib.gyro[1] >> 8);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS + 3, mpuCalib.gyro[1] & 0xff);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS + 4, mpuCalib.gyro[2] >> 8);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_BIAS + 5, mpuCalib.gyro[2] & 0xff);
	 */
}

static void motion_init() {
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_POWER, VALUE_POWER_RESET);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_POWER, VALUE_POWER_OFF);

	for (uint8_t i = 0; i < 5; i++) {
		uint8_t out = i2c_read_reg(hwnd, MPU_I2C_ADDR, REG_WHO_AM_I);
		if (out == VALUE_WHO_AM_I) {
			break;
		}
		if (i > 3) {
			return;
		}
	}

	HAL_Delay(100);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_MASTER_I2C_CTRL, 0x00);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_I2C_PASSTHROUGH, VALUE_I2C_PASSTHROUGH);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_POWER_2, VALUE_POWER_2_ENABLE_ALL);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_SAMPLE_RATE, VALUE_SAMPLE_RATE_NONE);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_CONFIG, VALUE_CONFIG_LPF_BW_BASED);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_GYRO_CONFIG, VALUE_GYRO_CONFIG_USE_HP_ADVICE);
	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_ACC_CONFIG, VALUE_ACC_CONFIG_4G);
	//i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_USER, VALUE_USER_ENABLE_I2C);

	// mag init
	for (uint8_t i = 0; i < 5; i++) {
		uint8_t id = mag_read_reg(REG_MAG_ID);
		if (id == VALUE_MAG_ID) {
			break;
		}
		if (i > 3) {
			return;
		}
	}

	mag_write_reg(REG_MAG_CTRL_2, 1);
	mag_write_reg(REG_MAG_CTRL, 0);
	mag_write_reg(REG_MAG_CTRL, 0x0f);
	mag_read(buf, REG_MAG_ASA_X, 3);
	mpuCalib.magneticScaling[0] = (float) (buf[0] - 128) / 256.f + 1.f;
	mpuCalib.magneticScaling[1] = (float) (buf[1] - 128) / 256.f + 1.f;
	mpuCalib.magneticScaling[2] = (float) (buf[2] - 128) / 256.f + 1.f;
	mag_write_reg(REG_MAG_CTRL, 0x11);

	i2c_write_reg(hwnd, MPU_I2C_ADDR, REG_POWER, VALUE_POWER_ON);
	if (system_is_module_ready(MODULE_STORAGE)) {
		// TODO: load stuff
	}

	motion_calibrate();

	system_mark_module_ready(MODULE_GYROSCOPE);
}

static void barometer_tick() {
	barometer_read_temperature();
	barometer_read_pressure();
}

static int16_t updateXqueue(int16_t x) {
	aXq[pXq] = x;
	if (++pXq >= 5) {
		pXq = 0;
	}
	int sum = 0;
	for (uint8_t i = 0; i < 5; i++) {
		sum += aXq[i];
	}
	return sum / 5;
}

static int16_t updateYqueue(int16_t x) {
	aYq[pYq] = x;
	if (++pYq >= 5) {
		pYq = 0;
	}
	int sum = 0;
	for (uint8_t i = 0; i < 5; i++) {
		sum += aYq[i];
	}
	return sum / 5;
}

static int16_t updateZqueue(int16_t x) {
	aZq[pZq] = x;
	if (++pZq >= 5) {
		pZq = 0;
	}
	int sum = 0;
	for (uint8_t i = 0; i < 5; i++) {
		sum += aZq[i];
	}
	return sum / 5;
}

static void motion_tick() {
	uint32_t now = HAL_GetTick();
	if (now - lastUpdateTime < 10) {
		return;
	}
	i2c_read(hwnd, MPU_I2C_ADDR, REG_ACC_DATA, buf, 6);
	lastMpuFrame.acceleration[X_AXIS] = updateXqueue(s16be(buf));
	lastMpuFrame.acceleration[Y_AXIS] = updateYqueue(s16be(buf + 2));
	lastMpuFrame.acceleration[Z_AXIS] = updateZqueue(s16be(buf + 4));
	i2c_read(hwnd, MPU_I2C_ADDR, REG_GYRO_DATA, buf, 6);
	lastMpuFrame.gyro[X_AXIS] = s16be(buf) - mpuCalib.gyro[X_AXIS];
	lastMpuFrame.gyro[Y_AXIS] = s16be(buf + 2) - mpuCalib.gyro[Y_AXIS];
	lastMpuFrame.gyro[Z_AXIS] = s16be(buf + 4) - mpuCalib.gyro[Z_AXIS];
	if (now - lastMagUpdateTime > 20) {
		uint8_t magReady = mag_read_reg(REG_MAG_STATE) & 0x01;
		if (magReady) {
			mag_read(buf, REG_MAG_DATA, 6);
			uint8_t magFailed = mag_read_reg(REG_MAG_STATE_2) & 0x08, magUpdate = 0;
			if (!magFailed) {
				lastMpuFrame.magnetic[0] = s16le(buf);
				lastMpuFrame.magnetic[1] = s16le(buf + 2);
				lastMpuFrame.magnetic[2] = s16le(buf + 4);
				lastMagUpdateTime = now;
				for (uint8_t i = 0; i < 3; i++) {
					if (lastMpuFrame.magnetic[i] > mpuCalib.magneticUpperBound[i]) {
						mpuCalib.magneticUpperBound[i] = lastMpuFrame.magnetic[i];
					}
					if (lastMpuFrame.magnetic[i] < mpuCalib.magneticLowerBound[i]) {
						mpuCalib.magneticLowerBound[i] = lastMpuFrame.magnetic[i];
					}
					if (mpuCalib.magneticLowerBound[i] < mpuCalib.magneticUpperBound[i]) {
						mpuCalib.magneticNorm[i] = (mpuCalib.magneticUpperBound[i] - mpuCalib.magneticLowerBound[i]) / 2;
						mpuCalib.magneticBias[i] = mpuCalib.magneticUpperBound[i] - mpuCalib.magneticNorm[i];
					}
				}
			}
			mag_write_reg(REG_MAG_CTRL, 0x11);
		}
	}
	lastMpuFrame.lastUpdateTime = lastUpdateTime;
}

void mpu_init(I2C_HandleTypeDef *handle) {
	hwnd = handle;
	motion_init();
	//barometer_init();
}

void mpu_tick() {
	if (system_is_module_ready(MODULE_GYROSCOPE)) {
		motion_tick();
	}
	if (system_is_module_ready(MODULE_BAROMETER)) {
		barometer_tick();
	}
}

const MpuData *mpu_get() {
	return &lastMpuFrame;
}

const BarometerData *baro_get() {
	return &lastBarometerFrame;
}

const MpuCalibrationData *mpu_calibration_get() {
	return &mpuCalib;
}
